// Copyright (C) 2010 and onwards Noam Shazeer and Georges Harik
// Pieces of this code may be licenced under the Apache License, version 2.0 
//   from Harik Shazeer Labs

#ifndef _RESTAURANT_H_
#define _RESTAURANT_H_

#include "util.h"
#include "table-structures.h"
#include "nolock.h"

inline int64 SelfModelingBooleanDL(uint64 a, uint64 b) {
  return MLnF(a+b+1) - MLnF(a) - MLnF(b); }

inline int64 PermutationDL(uint64 total, uint64 used) {
  CHECK(total >= used);
  return MLnF(total) - MLnF(total - used);
}

struct ThetaDist {
  vector<uint64> thetas_;
  vector<double> priors_;
  vector<double> ln_priors_;
  vector<double> ln_theta_;
  vector<double> ln_gamma_theta_;
  NLCache * factor_cache_;

  ThetaDist() { factor_cache_ = NULL; } 
  ~ThetaDist() { delete [] factor_cache_; }
  
  void Init(const vector<uint64>& thetas, const vector<double>& priors);
  void Init(uint64 levels, uint64 base, bool use_zero);
  
  int64 Factor(uint64 num_tables, uint64 total_count) const;
};

extern ThetaDist * G_THETAS;

void InitRestaurant();

// returns the ln of a factor in the probability of a partition under
// the generalized chinese restaurant process. 
// in particular 
//    ln[ GAMMA(theta) * theta ^ num_tables / GAMMA(theta + total_count) ] 
// The other factor in the partition probability is the product of the
//   gamma functions of all the table sizes. 
double RestaurantFactor(uint64 num_tables, 
			uint64 total_count, 
			uint64 theta);

// Count 0, 1, B, B^2, ... B^L
const uint64 kRestaurantLevels = 4;

int64 RestaurantUpdate(uint64 old_count, uint64 new_count, 
		       uint64 * total_count, uint64 *num_tables, 
		       const ThetaDist & thetas);

int64 ImmutableRestaurantUpdate
(uint64 old_count, uint64 new_count, 
 uint64 old_total_count, uint64 old_num_tables, 
 const ThetaDist & thetas);

int64 DinerUpdate (Table<Thing> * h, uint64 diner, uint64 key, int64 diff, int64 dlkey);

// a helper for computing diffs in the restaurant DL. 
// accounts for the factor plus the one table that is changing. 
int64 PartialRestaurantDL(uint64 count, uint64 total_count, uint64 num_tables,
			  const ThetaDist &thetas);


#endif

